Skip to content

Fix: stabilize test after distributed training completion #375

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Fix: stabilize test after distributed training completion #375

wants to merge 1 commit into from
+3 −1

Conversation

@stop1one
Copy link

@stop1one stop1one commented Oct 1, 2025

Description

Related Issues

Fixes #316 #374

Summary of Changes

This PR fixes two problems for distributed training when run_test=True:

  1. Synchronization before testing

    • Added torch.distributed.barrier() when args.distributed is enabled, right after training and before testing.
    • Ensures all processes wait until rank 0 has finished saving checkpoint_best_total.pth before moving on to testing.

  2. Correct checkpoint loading in DDP

    • Changed from model.load_state_dict(best_state_dict) to model_without_ddp.load_state_dict(best_state_dict)
    • Prevents key mismatch errors when loading checkpoints in distributed environments.

Type of change

Please delete options that are not relevant.

  • Bug fix (non-breaking change which fixes an issue)

How has this change been tested, please provide a testcase or example of how you tested the change?

train.py :

from rfdetr import RFDETRLarge

model = RFDETRLarge()

model.train(
    dataset_dir=<DATASET_PATH>,
    epochs=30,
    batch_size=8
    grad_accum_steps=2,
    lr=1e-4,
    output_dir=<OUTPUT_PATH>,
)

python -m torch.distributed.launch --nproc_per_node=4 --use_env train.py

@CLAassistant
Copy link

CLAassistant commented Oct 1, 2025 •
edited
Loading

CLA assistant check
All committers have signed the CLA.

@stop1one
Copy link
Author

stop1one commented Nov 6, 2025

Hi @probicheaux ,
I just realized I hadn’t signed the CLA earlier — that’s done now ✅
Would you mind taking another look at this PR when you have a moment?
Thanks again for your time!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@SkalskiP SkalskiP Awaiting requested review from SkalskiP SkalskiP is a code owner

@probicheaux probicheaux Awaiting requested review from probicheaux probicheaux is a code owner

@isaacrob-roboflow isaacrob-roboflow Awaiting requested review from isaacrob-roboflow isaacrob-roboflow is a code owner

@Matvezy Matvezy Awaiting requested review from Matvezy Matvezy is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

Distributed Training Fails at End: FileNotFoundError and State Dict Mismatch Issues

2 participants

@stop1one @CLAassistant